using System;
using System.Collections.Generic;
using System.Drawing;
using System.Drawing.Imaging;
using System.Text;

// The motion detector can establish a baseline for the background
// of a scene and then makes some predictions about what
// regions of an image contain relevant foreground objects

namespace Blimpbots.Recognition
{
    public class MotionDetector
    {
        Random r = new Random();

        byte[] _shifts = null;
        byte[] _objectMask = null;
        byte[] _bkgSamples = null;

        Bitmap _sceneAvg = null;
        Bitmap _scratchPad = null;
        Graphics _scratchG = null;

        double _shiftMean = 0;
        double _shiftVariance = 0;

        double scaleX, scaleY;

        public MotionDetector(int width, int height)
        {
            _scratchPad = new Bitmap(width, height, PixelFormat.Format24bppRgb);
            _scratchG = Graphics.FromImage(_scratchPad);
            _sceneAvg = new Bitmap(_scratchPad.Width, _scratchPad.Height, System.Drawing.Imaging.PixelFormat.Format24bppRgb);
            _objectMask = new byte[width * height];
            _shifts = new byte[width * height];
            _bkgSamples = new byte[width * height];
        }

        public void ObserveObjects(Rectangle[] objects)
        {
            for (int i = 0; i < objects.Length; i++)
            {
                // Rescale the rectangles.  TODO: move this scaling to a helper func
                objects[i] = new Rectangle((int)(objects[i].Left / scaleX), (int)(objects[i].Top / scaleY),
                                           (int)(objects[i].Width / scaleX), (int)(objects[i].Height / scaleY));
            }

            // Using the last frame, update our background estimate
            // for all the regions where the rectangles don't intersect

            // TODO: maybe just alloc new memory here?
            Array.Clear(_objectMask, 0, _objectMask.Length);
            CreateMask(objects, _scratchPad.Width, _objectMask);

            // lock scene image
            BitmapData sceneData = _scratchPad.LockBits(
                new Rectangle(0, 0, _scratchPad.Width, _scratchPad.Height),
                ImageLockMode.ReadOnly, _scratchPad.PixelFormat);

            BitmapData bkgData = _sceneAvg.LockBits(
                new Rectangle(0, 0, _sceneAvg.Width, _sceneAvg.Height),
                ImageLockMode.ReadOnly, _sceneAvg.PixelFormat);

            MergeBackground(sceneData, _shifts, _objectMask,
                            _bkgSamples, bkgData, _shiftMean, _shiftVariance);

            _sceneAvg.UnlockBits(bkgData);
            _scratchPad.UnlockBits(sceneData);

        }

        public void CreateMask(Rectangle[] objects, int width, byte[] mask)
        {
            if (mask.Length % width != 0)
                throw new Exception("Mask size must be a multiple of width.");

            int height = mask.Length / width;

            // Apply the mask for every rectangle
            foreach (Rectangle r in objects)
            {
                int bottom = Math.Min(r.Bottom, height);
                int right = Math.Min(r.Right, width);
                int left = Math.Max(r.Left, 0);
                int top = Math.Max(r.Top, 0);

                // Apply the rectangle
                for (int y = top; y < bottom; y++)
                {
                    int idx = left + y * width;
                    for (int x = left; x < right; x++, idx++)
                    { mask[idx] = 1; }
                }
            }
        }


        public Rectangle[] PredictForeground(Bitmap scene)
        {
            scaleX = (double)scene.Width / _sceneAvg.Width;
            scaleY = (double)scene.Height / _sceneAvg.Height;

            // Draw a downsampled copy of the image
            _scratchG.DrawImage(scene, 0, 0, _scratchPad.Width, _scratchPad.Height);

            // lock scene image
            BitmapData sceneData = _scratchPad.LockBits(
                new Rectangle(0, 0, _scratchPad.Width, _scratchPad.Height),
                ImageLockMode.ReadOnly, _scratchPad.PixelFormat);

            BitmapData bkgData = _sceneAvg.LockBits(
                new Rectangle(0, 0, _sceneAvg.Width, _sceneAvg.Height),
                ImageLockMode.ReadOnly, _sceneAvg.PixelFormat);

            double avgShift = GetShifts(sceneData, _shifts, bkgData);
            double shiftVar = EstimateShiftVariance(_shifts, avgShift, 50);

            /* Update statistic estimates */
            _shiftMean = 0.95 * _shiftMean + 0.05 * avgShift;
            _shiftVariance = 0.95 * _shiftVariance + 0.05 * shiftVar;

            _sceneAvg.UnlockBits(bkgData);
            _scratchPad.UnlockBits(sceneData);

            return GetObjects();
        }

        private Rectangle[] GetObjects()
        {
            Rectangle[] objectRects = FindObjects(_shifts, _shiftMean, _shiftVariance);

            // Rescale the rectangles
            for (int i = 0; i < objectRects.Length; i++)
            {
                Rectangle r = objectRects[i];
                objectRects[i] = new Rectangle((int)(r.Left * scaleX), (int)(r.Top * scaleY),
                                               (int)(r.Width * scaleX), (int)(r.Height * scaleY));
            }

            return objectRects;
        }

        public Bitmap DrawForeground()
        {
            Bitmap fg = _scratchPad;

            // lock scene image
            BitmapData fgData = fg.LockBits(
                new Rectangle(0, 0, fg.Width, fg.Height),
                ImageLockMode.ReadOnly, fg.PixelFormat);

            HideBackground(fgData, _shifts, _shiftMean, _shiftVariance);

            fg.UnlockBits(fgData);

            return fg;
        }

        public Bitmap GetBackground()
        {
            return _sceneAvg;
        }

        /// <summary>
        /// Quick monte-carlo estimate of the variance
        /// </summary>
        private double EstimateShiftVariance(byte[] shiftArray, double mean, int samples)
        {
            if (samples <= 1) return 0;
            double sqErrSum = 0;
            for(int i = 0; i < samples; i++)
                sqErrSum += Math.Pow((double)shiftArray[r.Next(shiftArray.Length)] - mean, 2);
            return sqErrSum / (samples - 1);
        }

        private Rectangle[] FindObjects(byte[] shiftArray, double shiftMean, double shiftVar)
        {
            int BOXES_X = 8;
            int BOXES_Y = 8;

            int width = _sceneAvg.Width;
            int height = _sceneAvg.Height;

            int box_w = width / BOXES_X;
            int box_h = height / BOXES_Y;

            int[,] boxes = new int[BOXES_X, BOXES_Y];

            int box_y, box_x, pix_idx = 0;

            // Pixel is "on" if it is more than 0.75 st dev from the mean
            int thresh = (int)(shiftMean + 0.75*Math.Sqrt(shiftVar));

            // for each line
            for (int y = 0; y < height; y++)
            {
                box_y = y / box_h;
                // for each pixel
                for (int x = 0; x < width; x++)
                {
                    box_x = x / box_w;
                    boxes[box_x, box_y] += (shiftArray[pix_idx] >= thresh) ? 1 : 0;
                    pix_idx++;
                }
            }

            // Box contains motion if 1/4 of its pixels are moving
            // TODO: this is a tuning parameter, as is the previous
            //       threshold
            int boxThresh = (box_h * box_w) / 4;

            LinkedList<Rectangle> objectRects = new LinkedList<Rectangle>();

            for (box_x = 0; box_x < BOXES_X; box_x++)
                for (box_y = 0; box_y < BOXES_Y; box_y++)
                {
                    if (boxes[box_x, box_y] > boxThresh)
                    {
                        Rectangle r = new Rectangle(box_w * box_x, box_h * box_y, box_w, box_h);
                        r.Inflate(box_w / 2, box_h / 2);
                        objectRects.AddLast(r);
                    }
                }

            Rectangle[] objectRectArr = Blimpbots.Recognition.Euclidian.Rectangles.Merge(objectRects);

            return objectRectArr;
        }

        private unsafe void MergeBackground(BitmapData fgData, byte[] shiftArray, byte[] objectMask, byte[] bkgSamples, BitmapData bkgData, double shiftMean, double shiftVar)
        {
            // get image dimension
            int width = fgData.Width;
            int height = fgData.Height;

            // initialize other variables
            int pixelSize = 3;
            int lineSize = width * pixelSize;
            int offset = fgData.Stride - lineSize;
            // pixel value
            int v;

            // do the job
            byte* ptr = (byte*)fgData.Scan0.ToPointer();
            byte* ovr = (byte*)bkgData.Scan0.ToPointer();

            int pix_idx = 0; // overall pixel index

            double wt = 0; int wt_sum = 0;

            // for each line
            for (int y = 0; y < height; y++)
            {
                // for each pixel
                for (int x = 0; x < width; x++)
                {
                    // determine how much to weight this pixel into the bkg                    
                    wt = (double)(shiftMean - shiftArray[pix_idx]) / 255;
                    wt = (wt < 0) ? 2 : (wt * 255) + 2;

                    wt_sum = bkgSamples[pix_idx] + (int)wt;

                    for (int color = 0; color < 3; color++)
                    {
                        if (bkgSamples[pix_idx] > 10)
                        {
                            v = (*ptr > *ovr) ? *ovr + 1 : *ovr - 1;
                            if (v < 0) v = 0;
                            else if (v > 255) v = 255;
                        }
                        else
                        {
                            v = bkgSamples[pix_idx] * (int)*ovr + (int)wt * (int)*ptr;
                            v /= wt_sum;
                        }

                        if (objectMask[pix_idx] == 0)
                            *ovr = (byte)v; // adjust the bkg

                        ptr++; ovr++;
                    }

                    if(objectMask[pix_idx] == 0)
                        if (bkgSamples[pix_idx] < 255)
                            bkgSamples[pix_idx]++;

                    pix_idx++; // go to next pixel
                }

                ptr += offset;
                ovr += offset;
            }

        }

        private unsafe void HideBackground(BitmapData fgData, byte[] shiftArray, double shiftMean, double shiftVar)
        {
            // get image dimension
            int width = fgData.Width;
            int height = fgData.Height;

            // initialize other variables
            int pixelSize = 3;
            int lineSize = width * pixelSize;
            int offset = fgData.Stride - lineSize;
            int brightness;

            // do the job
            byte* ptr = (byte*)fgData.Scan0.ToPointer();

            int pix_idx = 0; // overall pixel index

            // for each line
            for (int y = 0; y < height; y++)
            {
                // for each pixel
                for (int x = 0; x < width; x++)
                {
                    double shift = shiftArray[pix_idx] - shiftMean;
                    if (shift > shiftVar) brightness = 255;
                    else brightness = 75;

                    *ptr = (byte)((*ptr * brightness) / 255); ptr++;
                    *ptr = (byte)((*ptr * brightness) / 255); ptr++;
                    *ptr = (byte)((*ptr * brightness) / 255); ptr++;

                    pix_idx++;
                }
                ptr += offset;
            }

        }

        private unsafe double GetShifts(BitmapData fgData, byte[] shiftArray, BitmapData bkgData)
        {
            long shiftSum = 0;
            int shiftCount = 0;

            // get image dimension
            int width = fgData.Width;
            int height = fgData.Height;

            // initialize other variables
            int pixelSize = 3;
            int lineSize = width * pixelSize;
            int offset = fgData.Stride - lineSize;
            // pixel value
            int v, sumdiff;

            // do the job
            byte* ptr = (byte*)fgData.Scan0.ToPointer();
            byte* ovr = (byte*)bkgData.Scan0.ToPointer();

            int pix_idx = 0; // overall pixel index

            #region Compute shifts over all pixels

            // for each line
            for (int y = 0; y < height; y++)
            {
                // for each pixel
                for (int x = 0; x < width; x++)
                {
                    sumdiff = 0;
                    for (int color = 0; color < 3; color++)
                    {
                        v = (int)*ptr - (int)*ovr;
                        // sumdiff += v * v;// (v < 0) ? -v : v;
                        sumdiff += (v < 0) ? -v : v;
                        ptr++; ovr++;
                    }

                    // euclidian distance makes for smoother bkg incorporation
                    //sumdiff = (int)Math.Sqrt(sumdiff);

                    shiftArray[pix_idx] = (sumdiff > 255) ? (byte)255 : (byte)sumdiff;

                    shiftSum += shiftArray[pix_idx]; shiftCount += 1;

                    pix_idx++; // go to next pixel
                }

                ptr += offset;
                ovr += offset;
            }

            #endregion

            // Return the mean shift value
            return (shiftCount == 0) ? 0 : (double)shiftSum / shiftCount;
        }

    } // end class

} // end namespace
